# coding = utf-8

'''
m-Unet的代码
'''

import torch
import torch.nn as nn


def double_conv(in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        #nn.Dropout2d(0.1),
        nn.Conv2d(out_channels, out_channels, 3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        #nn.Dropout2d(0.1)
    )

def object_filter(in_channels, out_channels):
    return nn.Sequential(
        nn.ConvTranspose2d(in_channels, out_channels, kernel_size=2, stride=2),
        #nn.Upsample(scale_factor=2, mode="bilinear")
        nn.ReLU(inplace=True)
    )

def skip_conv(in_channels, out_channels):
    return nn.Sequential(
        nn.Conv2d(in_channels, out_channels, 3, padding=1),
        nn.BatchNorm2d(out_channels),
        nn.ReLU(inplace=True),
        #nn.Dropout2d(0.1)

    )



class MUNet(nn.Module):

    def __init__(self, input_channel, n_class):
        super().__init__()

        self.dconv_down1 = double_conv(input_channel, 64)
        self.dconv_down2 = double_conv(64, 128)
        self.dconv_down3 = double_conv(128, 256)
        self.dconv_down4 = double_conv(256, 512)
        self.dconv_down5 = double_conv(512, 1024)

        self.filter_down1 = object_filter(64, 64)
        self.filter_down2 = object_filter(128, 128)
        self.filter_down3 = object_filter(256, 256)
        self.filter_down4 = object_filter(512, 512)

        self.skip_down1 = skip_conv(64, 64)
        self.skip_down2 = skip_conv(128, 128)
        self.skip_down3 = skip_conv(256, 256)
        self.skip_down4 = skip_conv(512, 512)


        self.maxpool = nn.MaxPool2d(2)
        #self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.up_sample4 = nn.ConvTranspose2d(1024, 1024, stride=2, kernel_size=2)
        self.up_sample3 = nn.ConvTranspose2d(512, 512, stride=2, kernel_size=2)
        self.up_sample2 = nn.ConvTranspose2d(256, 256, stride=2, kernel_size=2)
        self.up_sample1 = nn.ConvTranspose2d(128, 128, stride=2, kernel_size=2)


        self.dconv_up4 = double_conv(1024 + 512, 512)
        self.dconv_up3 = double_conv(256 + 512, 256)
        self.dconv_up2 = double_conv(128 + 256, 128)
        self.dconv_up1 = double_conv(128 + 64, 64)

        self.conv_last = nn.Conv2d(64, n_class, 1)

    def forward(self, x):
        conv1 = self.dconv_down1(x)
        x = self.maxpool(conv1)
        filter1 = self.filter_down1(x)
        skip1 = self.skip_down1(conv1 - filter1)

        conv2 = self.dconv_down2(x)
        x = self.maxpool(conv2)
        filter2 = self.filter_down2(x)
        skip2 = self.skip_down2(conv2 - filter2)

        conv3 = self.dconv_down3(x)
        x = self.maxpool(conv3)
        filter3 = self.filter_down3(x)
        skip3 = self.skip_down3(conv3 - filter3)

        conv4 = self.dconv_down4(x)
        x = self.maxpool(conv4)
        filter4 = self.filter_down4(x)
        skip4 = self.skip_down4(conv4 - filter4)

        x = self.dconv_down5(x)

        x = self.up_sample4(x)
        x = torch.cat([x, skip4], dim=1)
        x = self.dconv_up4(x)

        x = self.up_sample3(x)
        x = torch.cat([x, skip3], dim=1)
        x = self.dconv_up3(x)

        x = self.up_sample2(x)
        x = torch.cat([x, skip2], dim=1)
        x = self.dconv_up2(x)

        x = self.up_sample1(x)
        x = torch.cat([x, skip1], dim=1)
        x = self.dconv_up1(x)


        out = self.conv_last(x)

        return out

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                m.weight.data.normal_(0, 0.05)
                if m.bias is not None:
                    m.bias.data.fill_(0.1)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

class UNetDeep(nn.Module):

    def __init__(self, input_channel, n_class):
        super().__init__()

        self.dconv_down1 = double_conv(input_channel, 64)
        self.dconv_down2 = double_conv(64, 128)
        self.dconv_down3 = double_conv(128, 256)
        self.dconv_down4 = double_conv(256, 512)
        self.dconv_down5 = double_conv(512, 1024)
        self.dconv_down6 = double_conv(1024, 2048)
        self.dconv_down7 = double_conv(2048, 5096)

        self.filter_down1 = object_filter(64, 64)
        self.filter_down2 = object_filter(128, 128)
        self.filter_down3 = object_filter(256, 256)
        self.filter_down4 = object_filter(512, 512)
        self.filter_down5 = object_filter(1024, 1024)
        self.filter_down6 = object_filter(2048, 2048)

        self.skip_down1 = skip_conv(64, 64)
        self.skip_down2 = skip_conv(128, 128)
        self.skip_down3 = skip_conv(256, 256)
        self.skip_down4 = skip_conv(512, 512)
        self.skip_down5 = skip_conv(1024, 1024)
        self.skip_down6 = skip_conv(2048, 2048)


        self.maxpool = nn.MaxPool2d(2)
        #self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
        self.up_sample6 = nn.ConvTranspose2d(5096, 5096, stride=2, kernel_size=2)
        self.up_sample5 = nn.ConvTranspose2d(2048, 2048, stride=2, kernel_size=2)
        self.up_sample4 = nn.ConvTranspose2d(1024, 1024, stride=2, kernel_size=2)
        self.up_sample3 = nn.ConvTranspose2d(512, 512, stride=2, kernel_size=2)
        self.up_sample2 = nn.ConvTranspose2d(256, 256, stride=2, kernel_size=2)
        self.up_sample1 = nn.ConvTranspose2d(128, 128, stride=2, kernel_size=2)


        self.dconv_up6 = double_conv(5096 + 2048, 2048)
        self.dconv_up5 = double_conv(2048 + 1024, 1024)
        self.dconv_up4 = double_conv(1024 + 512, 512)
        self.dconv_up3 = double_conv(256 + 512, 256)
        self.dconv_up2 = double_conv(128 + 256, 128)
        self.dconv_up1 = double_conv(128 + 64, 64)

        self.conv_last = nn.Conv2d(64, n_class, 1)

    def forward(self, x):
        conv1 = self.dconv_down1(x)
        x = self.maxpool(conv1)
        filter1 = self.filter_down1(x)
        skip1 = self.skip_down1(conv1 - filter1)

        conv2 = self.dconv_down2(x)
        x = self.maxpool(conv2)
        filter2 = self.filter_down2(x)
        skip2 = self.skip_down2(conv2 - filter2)

        conv3 = self.dconv_down3(x)
        x = self.maxpool(conv3)
        filter3 = self.filter_down3(x)
        skip3 = self.skip_down3(conv3 - filter3)

        conv4 = self.dconv_down4(x)
        x = self.maxpool(conv4)
        filter4 = self.filter_down4(x)
        skip4 = self.skip_down4(conv4 - filter4)

        conv5 = self.dconv_down5(x)
        x = self.maxpool(conv5)
        filter5 = self.filter_down5(x)
        skip5 = self.skip_down5(conv5 - filter5)

        conv6 = self.dconv_down6(x)
        x = self.maxpool(conv6)
        filter6 = self.filter_down6(x)
        skip6 = self.skip_down6(conv6 - filter6)

        x = self.dconv_down7(x)

        x = self.up_sample6(x)
        x = torch.cat([x, skip6], dim=1)
        x = self.dconv_up6(x)

        x = self.up_sample5(x)
        x = torch.cat([x, skip5], dim=1)
        x = self.dconv_up5(x)

        x = self.up_sample4(x)
        x = torch.cat([x, skip4], dim=1)
        x = self.dconv_up4(x)

        x = self.up_sample3(x)
        x = torch.cat([x, skip3], dim=1)
        x = self.dconv_up3(x)

        x = self.up_sample2(x)
        x = torch.cat([x, skip2], dim=1)
        x = self.dconv_up2(x)

        x = self.up_sample1(x)
        x = torch.cat([x, skip1], dim=1)
        x = self.dconv_up1(x)


        out = self.conv_last(x)

        return out

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
                m.weight.data.normal_(0, 0.05)
                if m.bias is not None:
                    m.bias.data.fill_(0.1)
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                m.weight.data.normal_(0, 0.01)
                m.bias.data.zero_()

if __name__ == '__main__':
    unet = MUNet(input_channel=1, n_class=2).cuda()
    from torchsummary import summary

    summary(unet, (1, 224, 224))

    unet._initialize_weights()